import base64

from ..common import VolcengineClient
from ..util import convert_bytes_to_image, convert_image_to_bytes


class DouBaoFaceSwapAI:
    RETURN_TYPES = ("IMAGE",)
    # RETURN_NAMES = ("image_output_name",)

    FUNCTION = "image_creator"

    # OUTPUT_NODE = False

    CATEGORY = "DouBaoFaceSwapAI"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "info": ("STRING", {
                    "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                    "default": "{}",
                    "lazy": True
                }),
                "image": ("IMAGE",),
                "reference": ("IMAGE",),
                "gpen": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.1}),
                "skin": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.1}),
            },
        }

    @staticmethod
    def image_creator(info, image, reference, gpen: float = 0.9, skin: float = 0.9):
        image_numpy = image
        if isinstance(image, list):
            for img in image:
                image_numpy = img
        else:
            image_numpy = image

        reference_numpy = reference
        if isinstance(image, list):
            for img in image:
                reference_numpy = img
        else:
            reference_numpy = reference
        image_bytes = convert_image_to_bytes(image_numpy)
        image_base64 = base64.b64encode(image_bytes).decode("utf-8")
        reference_bytes = convert_image_to_bytes(reference_numpy)
        reference_base64 = base64.b64encode(reference_bytes).decode("utf-8")
        client = VolcengineClient(volcengine_info=info)
        operate_result = client.faceswap_ai(image_base64s=[image_base64, reference_base64], gpen=gpen,
                                            skin=skin)
        if operate_result.code == 1:
            pics = operate_result.data
            image_bytes = pics[0]
            image = convert_bytes_to_image(image_bytes)
        return (image,)
